package com.czh.subject.infra.basic.es;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import io.micrometer.common.util.StringUtils;
import jakarta.annotation.PostConstruct;
import lombok.extern.slf4j.Slf4j;
import org.apache.http.HttpHost;
import org.apache.http.util.EntityUtils;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.delete.DeleteRequest;
import org.elasticsearch.action.delete.DeleteResponse;
import org.elasticsearch.action.get.GetRequest;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.search.SearchType;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.client.*;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.reindex.BulkByScrollResponse;
import org.elasticsearch.index.reindex.DeleteByQueryRequest;
import org.elasticsearch.index.reindex.UpdateByQueryRequest;
import org.elasticsearch.script.Script;
import org.elasticsearch.search.Scroll;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
import org.elasticsearch.search.sort.SortOrder;
import org.springframework.stereotype.Component;

import java.util.*;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * elasticsearch 操作客户端
 *
 * @author caozaihui
 * @version 1.0
 * @since 2024/9/9 22:03
 */
@Component
@Slf4j
public class EsRestClient {

    private static final Map<String, RestHighLevelClient> restHighLevelClientMap = new LinkedHashMap<>();

    private final EsConfigProperties esConfigProperties;

    private final RequestOptions requestOptions ;

    public EsRestClient(EsConfigProperties esConfigProperties){
        this.esConfigProperties = esConfigProperties;
        requestOptions = RequestOptions.DEFAULT.toBuilder().build();
    }

    @PostConstruct
    public void init() {
        for (EsClusterConfig esClusterConfig : esConfigProperties.getEsClusterConfigList()) {
            RestHighLevelClient restHighLevelClient = initRestClient(esClusterConfig);
            restHighLevelClientMap.put(esClusterConfig.getName(), restHighLevelClient);
        }
        log.info("\n=====================当前{}个ESClient 初始化成功！=====================\n", restHighLevelClientMap.size());
    }

    private RestHighLevelClient initRestClient(EsClusterConfig esClusterConfig) {
        List<HttpHost> httpHostList = new LinkedList<>();
        for (String ipPort : esClusterConfig.getNodes().split(",")) {
            String[] ipPortInfo = ipPort.split(":");
            if (ipPortInfo.length == 2) {
                HttpHost httpHost = new HttpHost(ipPortInfo[0], Integer.parseInt(ipPortInfo[1]));
                httpHostList.add(httpHost);
            }
        }
        HttpHost[] httpHostArray = new HttpHost[httpHostList.size()];
        httpHostList.toArray(httpHostArray);

        RestClientBuilder builder = RestClient.builder(httpHostArray);
        return new RestHighLevelClient(builder);
    }
    /**
     * 根据restClientName获取对应客户端
     *
     * @param restClientName 客户端名称
     * @return 客户端
     */
    private RestHighLevelClient getRestClient(String restClientName) {
        return restHighLevelClientMap.get(restClientName);
    }

    /**
     * 插入文档
     */
    public boolean insertDoc(EsIndexInfo esIndexInfo, EsSourceData esSourceData) {
        try {
            IndexRequest indexRequest = new IndexRequest(esIndexInfo.getIndexName()).id(esSourceData.getDocId()).source(esSourceData.getData());
            getRestClient(esIndexInfo.getClusterName()).index(indexRequest, requestOptions);
            log.info("-->esDoc插入{}成功!", JSON.toJSONString(indexRequest));
        } catch (Exception e) {
            log.error("insertDoc.exception:{}", e.getMessage(), e);
            return false;
        }
        return true;
    }

    /**
     * 删除所有
     */
    public boolean delete(EsIndexInfo esIndexInfo) {
        DeleteByQueryRequest deleteByQueryRequest = new DeleteByQueryRequest(esIndexInfo.getIndexName());
        deleteByQueryRequest.setQuery(QueryBuilders.matchAllQuery());
        try {
            BulkByScrollResponse bulkByScrollResponse = getRestClient(esIndexInfo.getClusterName()).deleteByQuery(deleteByQueryRequest, requestOptions);
            long deleted = bulkByScrollResponse.getDeleted();
            log.info("-->删除全部esDoc:{}", deleted);
            return deleted > 0;
        } catch (Exception e) {
            log.error("delete.exception:{}", e.getMessage(), e);
            return false;
        }
    }

    /**
     * 根据id删除文档
     */
    public boolean deleteDoc(EsIndexInfo esIndexInfo, String docId) {
//        DeleteByQueryRequest deleteByQueryRequest = new DeleteByQueryRequest(esIndexInfo.getIndexName());
//        deleteByQueryRequest.setQuery(QueryBuilders.matchQuery("id", docId));
        DeleteRequest deleteRequest = new DeleteRequest(esIndexInfo.getIndexName()).id(docId);
        try {
            DeleteResponse response = getRestClient(esIndexInfo.getClusterName()).delete(deleteRequest, requestOptions);
            log.info("-->删除esDoc:{}", JSON.toJSONString(response));
            return true;
        } catch (Exception e) {
            log.error("deleteDoc.exception:{}", e.getMessage(), e);
            return false;
        }
    }

    /**
     * 查询对应id文档是否存在
     */
    public boolean isExitsDocById(EsIndexInfo esIndexInfo, String docId) {
        GetRequest getRequest = new GetRequest(esIndexInfo.getIndexName()).id(docId);
        try {
            return getRestClient(esIndexInfo.getClusterName()).exists(getRequest, requestOptions);
        } catch (Exception e) {
            log.error("isExitsDocById.exception:{}", e.getMessage(), e);
            return false;
        }
    }

    /**
     * 根据id查询文档
     */
    public Map<String, Object> getDocById(EsIndexInfo esIndexInfo, String docId) {
        GetRequest getRequest = new GetRequest(esIndexInfo.getIndexName()).
                id(docId);
        try {
            return getRestClient(esIndexInfo.getClusterName()).get(getRequest, requestOptions).getSource();
        } catch (Exception e) {
            return Collections.emptyMap();
        }
    }

    /**
     * 根据id查询，显示指定字段
     *
     * @param docId  文档id
     * @param fields 需要展示的字段
     */
    public Map<String, Object> getDocById(EsIndexInfo esIndexInfo, String docId, String[] fields) {
        GetRequest request = new GetRequest(esIndexInfo.getIndexName(), docId);
        FetchSourceContext fetchSourceContext = new FetchSourceContext(true, fields, null);
        request.fetchSourceContext(fetchSourceContext);
        try {
            GetResponse documentFields = getRestClient(esIndexInfo.getClusterName()).get(request, requestOptions);
            log.info("-->getDocById:{}", JSON.toJSONString(documentFields.getSource()));
            return documentFields.getSource();
        } catch (Exception e) {
            log.error("getDocById.exception:{}", e.getMessage(), e);
            return Collections.emptyMap();
        }
    }
    /**
     * 使用Term Query进行搜索操作
     * 此方法主要用于在指定的Elasticsearch索引中执行搜索操作，支持高亮显示、排序以及滚动查询
     *
     * @param esIndexInfo     索引信息，包含索引名称和集群名称
     * @param esSearchRequest 搜索请求对象，携带了搜索的详细配置如查询条件、分页信息、排序信息等
     * @return 返回搜索响应对象，包含搜索结果等信息
     */
    public SearchResponse searchWithTermQuery(EsIndexInfo esIndexInfo, EsSearchRequest esSearchRequest) {
        BoolQueryBuilder queryBuilder = esSearchRequest.getBoolQueryBuilder();
        Long minutes = esSearchRequest.getMinutes();
        String[] fields = esSearchRequest.getFields();
        int size = esSearchRequest.getSize();
        int from = esSearchRequest.getFrom();
        Boolean needScroll = esSearchRequest.getNeedScroll();
        HighlightBuilder highlightBuilder = esSearchRequest.getHighlightBuilder();
        String sortName = esSearchRequest.getSortName();
        SortOrder sortOrder = esSearchRequest.getSortOrder();

        //搜索源构建器
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        searchSourceBuilder.query(queryBuilder);
        try {
            if (fields.length > 0){
                searchSourceBuilder.fetchSource(fields,null);
            }
            if (from>0){
                searchSourceBuilder.from(from);
            }
            if (size>0){
                searchSourceBuilder.size(size);
            }
            if (Objects.nonNull(highlightBuilder)) {
                searchSourceBuilder.highlighter(highlightBuilder);
            }
            if (StringUtils.isNotBlank(sortName)) {
                searchSourceBuilder.sort(sortName);
            } else if (StringUtils.isNotBlank(sortName)|| Objects.nonNull(sortOrder)) {
                searchSourceBuilder.sort(sortName,sortOrder);
            }
            // 初始化搜索请求
            SearchRequest searchRequest = new SearchRequest()
                    .searchType(SearchType.DEFAULT).indices(esIndexInfo.getIndexName()).source(searchSourceBuilder);
            if (needScroll) {
                Scroll scroll = new Scroll(TimeValue.timeValueMinutes(minutes));
                searchRequest.scroll(scroll);
            }
            return getRestClient(esIndexInfo.getClusterName()).search(searchRequest, requestOptions);
        }catch (Exception e){
            log.error("searchWithTermQuery.exception:{}", e.getMessage(), e);
            return null;
        }

    }
    /**
     * 更新文档
     * @param esSourceData 文档
     */
    public boolean updateDoc(EsIndexInfo esIndexInfo, EsSourceData esSourceData) {
        try {
            UpdateRequest updateRequest = new UpdateRequest(esIndexInfo.getIndexName(), esSourceData.getDocId())
                    .doc(esSourceData.getData());
            getRestClient(esIndexInfo.getClusterName()).update(updateRequest, requestOptions);
            log.info("-->updateDoc:{}", JSON.toJSONString(updateRequest));
        } catch (Exception e) {
            log.error("updateDoc.exception:{}", e.getMessage(), e);
            return false;
        }
        return true;
    }

    /**
     * 批量更新
     */
    public boolean batchUpdateDoc(EsIndexInfo esIndexInfo, List<EsSourceData> esSourceDataList) {
        if (log.isInfoEnabled()) {
            log.info("{}批量更新ES:{}条", esIndexInfo.getIndexName(), esSourceDataList.size());
        }
        if (esSourceDataList.isEmpty()) {
            return false;
        }
        AtomicBoolean temp = new AtomicBoolean(false);
        BulkRequest bulkRequest = new BulkRequest();
        esSourceDataList.forEach(esSourceData -> {
            UpdateRequest updateRequest = new UpdateRequest(esIndexInfo.getIndexName(), esSourceData.getDocId()).doc(esSourceData.getData());
            bulkRequest.add(updateRequest);
            temp.set(true);
        });
        try {
            if (temp.get()) {
                BulkResponse bulk = getRestClient(esIndexInfo.getClusterName()).bulk(bulkRequest, requestOptions);
                // 检查批量操作是否有失败
                log.info("-->bulk.batchUpdateDoc:{}", bulk.hasFailures());
                return !bulk.hasFailures();
            }
            return false;
        } catch (Exception e) {
            log.error("insertBatchDoc.exception:{}", e.getMessage(), e);
            return false;
        }
    }

    /**
     * 批量新增
     */
    public boolean batchInsertDoc(EsIndexInfo esIndexInfo, List<EsSourceData> esSourceDataList) {
        if (log.isInfoEnabled()) {
            log.info("{}批量新增ES:{}条", esIndexInfo.getIndexName(), esSourceDataList.size());
        }
        AtomicBoolean temp = new AtomicBoolean(false);
        BulkRequest bulkRequest = new BulkRequest();
        esSourceDataList.forEach(esSourceData -> {
            IndexRequest updateRequest = new IndexRequest(esIndexInfo.getClusterName()).id(esSourceData.getDocId()).source(esSourceData.getData());
            bulkRequest.add(updateRequest);
            temp.set(true);
        });
        try {
            if (temp.get()) {
                BulkResponse bulk = getRestClient(esIndexInfo.getClusterName()).bulk(bulkRequest, requestOptions);
                log.info("-->bulk.batchInsertDoc:{}", bulk.hasFailures());
                return !bulk.hasFailures();
            }
            return false;
        } catch (Exception e) {
            log.error("batchInsertDoc.exception:{}", e.getMessage(), e);
            return false;
        }
    }

    /**
     *使用查询更新Es索引数据
     * @param esIndexInfo 索引信息，包括索引名称等
     * @param queryBuilder 查询构建器，用于指定需要更新的文档
     * @param script 更新脚本，定义如何更新文档
     * @param batchSize 批处理大小，控制一次处理的文档数量
     * @return 如果更新操作成功，则返回true；否则返回false
     */
    public boolean updateByQuery(EsIndexInfo esIndexInfo, QueryBuilder queryBuilder, Script script, int batchSize) {
        if (log.isInfoEnabled()) {
            log.info("updateByQuery.indexName:{}", esIndexInfo.getIndexName());
        }
        try {
            UpdateByQueryRequest updateByQueryRequest = new UpdateByQueryRequest(esIndexInfo.getIndexName());
            updateByQueryRequest.setQuery(queryBuilder).setScript(script).setBatchSize(batchSize).setAbortOnVersionConflict(false);

            BulkByScrollResponse bulkByScrollResponse = getRestClient(esIndexInfo.getClusterName()).updateByQuery(updateByQueryRequest, requestOptions);
            log.info("-->updateByQuery:{}", bulkByScrollResponse.getBatches());
            return bulkByScrollResponse.getBatches()>0;
        }catch (Exception e){
            log.error("updateByQuery.exception:{}", e.getMessage(), e);
            return false;
        }
    }

    /**
     * 分词方法
     * @param esIndexInfo ES索引信息对象，包含集群名称等必要信息
     * @param text 待分词的文本字符串
     * @return 返回分词后的字符串列表
     */
    public List<String> getAnalyze(EsIndexInfo esIndexInfo, String text) throws Exception {
        List<String> list = new ArrayList<>();
        Request request = new Request("GET", "_analyze");
        // 创建请求体，指定分词器为ik_smart，并设置待分词文本
        JSONObject entity = new JSONObject();
        entity.put("analyzer", "ik_smart");
        entity.put("text", text);
        request.setJsonEntity(entity.toJSONString());
        Response response = getRestClient(esIndexInfo.getClusterName()).getLowLevelClient().performRequest(request);
        JSONObject tokens = JSONObject.parseObject(EntityUtils.toString(response.getEntity()));
        JSONArray arrays = tokens.getJSONArray("tokens");
        for (int i = 0; i < arrays.size(); i++) {
            JSONObject obj = JSON.parseObject(arrays.getString(i));
            list.add(obj.getString("token"));
        }
        log.info("-->getAnalyze:{}", list);
        return list;
    }
}